
from .vdo import VideoKeeper,VideoLoader
from .ldb import DataBase
import os,cv2


# 用于表示一个ReID的数据集中的一个类别

class ReidSubSet:

    def __init__(self,seq_db:DataBase,vdo:VideoLoader=None) -> None:
        self.seq_db = seq_db
        self.vdo = vdo
        self._classify()
        pass

    def get_crop_pitc(self,frameid_index_from_1:int,x,y,w,h):
        """
        从视频中获取指定的bbox位置的图像
        """
        if self.vdo is None:
            raise Exception("no video info inited to reidsubset")
        ret,img = self.vdo.getFrameImage(frameid_idx_from_0=frameid_index_from_1-1)
        if ret == False:
            print(ret)
            print(type(img))
            print(img)
            raise Exception(f"no frame crop found at frameid:{frameid_index_from_1},pos:{[x,y,w,h]}")
        return img[y:y+h,x:w+x]

    def _classify(self):
        """
        开始根据unqie_id进行分类
        """
        self.unique_name_dict = {}
        for row in self.seq_db:
            # unique_name = f"{self.seq_name}-{cname}-{inclasid}"
            unique_name,seq_name,frameid,inclasid,x,y,w,h,cname,cid = row
            info_tuple = ( seq_name,frameid,inclasid,x,y,w,h,cname,cid )
            self.unique_name_dict.setdefault(unique_name,[]).append(info_tuple)
        # 所有的unique类别
        self.unique_name_list = list(self.unique_name_dict.keys())
        self.unique_name_list.sort()
    
    def Unique_Name_List(self):
        return self.unique_name_list.copy()
    
    def __iter__(self):
        self.cur = 0
        self.lenth = len(self.unique_name_list)
        return self
    
    def __next__(self):
        if self.cur == self.lenth:
            raise StopIteration
        tuple_lis_key = self.unique_name_list[self.cur]
        self.cur += 1
        # 返回:unique_name,[(seq_name,frameid,inclasid,x,y,w,h,cname,cid)...]
        return tuple_lis_key,self.unique_name_dict[tuple_lis_key]




def path_mkdir_if_notexists(path):
    if os.path.exists(path) == False:
        os.mkdir(path)



class ReidDataSet:

    def __init__(self,root_path:str) -> None:
        self.unique_name_set = set()
        self.root_path = root_path
        path_mkdir_if_notexists(self.root_path)
        pass

    def AddSubSet(self,rst:ReidSubSet):
        for unique_name,lis in rst:
            if unique_name in self.unique_name_set:
                raise Exception(f"unique name {unique_name} has been here")
            # 记录该unique_name
            self.unique_name_set.add(unique_name)
            # 接下来将unique_name对应的文件夹生成
            pitc_path = os.path.join(self.root_path,unique_name)
            path_mkdir_if_notexists(pitc_path)
            # 获取图片并保存
            for seq_name,frameid,inclasid,x,y,w,h,cname,cid in lis:
                img = rst.get_crop_pitc(frameid,x,y,w,h)
                filename = f"{seq_name}-{frameid:0>5d}-{cname}-{inclasid}.jpg"
                file_path = os.path.join(pitc_path,filename)
                try:
                    cv2.imwrite(file_path,img)
                except Exception as e:
                    print(f"cv2.imwrite error img:{frameid,cname,inclasid} for size??={[x,y,w,h]}")
            print(f"build sub dataset for {unique_name} at {pitc_path} - ok")
    
    def GetAllUniqueNames(self):
        return self.unique_name_set.copy()

